//
// Created by guoshunw on 24-6-4.
//

#ifndef HELLO_SWAPCONTEXT_COROUTINE_H
#define HELLO_SWAPCONTEXT_COROUTINE_H

#include <ucontext.h>
#include <vector>
#include <functional>
#include <cassert>
#include <memory>

class Coroutine;

const static int STACK_SIZE = 1024 * 1024;

enum class CoroutineStatus {
    READY = 0,
    RUNNING,
    SUSPENDED,
    DEAD
};

class Schedule {
    char stack[STACK_SIZE]{};
    ucontext_t main{};
    std::unordered_map<long, std::shared_ptr<Coroutine>> coroutines{};
    long running = -1;

    friend class Coroutine;
    void mainFun(Coroutine *co);

public:
    explicit Schedule();

    bool hasCoroutine() { return !coroutines.empty(); };

    std::shared_ptr<Coroutine> addCoroutine(const std::function<int(Coroutine *)> &fun, const void *userData);

    std::shared_ptr<Coroutine> get(long id) {
        if (coroutines.find(id) == coroutines.end()) return {};
        return coroutines.at(id);
    };

    ~Schedule();
};


class Coroutine {
    std::function<int(Coroutine* userData)> fun{};
    const void *userData{};
    ucontext_t ctx{};
    Schedule *sch{};

    std::ptrdiff_t capacity = 0;
    std::ptrdiff_t size = 0;

public:
    [[nodiscard]] const void *getUserData() const;

public:
    [[nodiscard]] long getId() const;

private:
    static void mainFun(Schedule *schedule, Coroutine *co);

    void saveStack(const char *stack_top);

    CoroutineStatus status = CoroutineStatus::READY;
    void *stack{};

    static long getNextId() {
        static long nextId = -1;
        return ++nextId;
    }

    long id = -1;

public:
    explicit Coroutine(Schedule *schedule, const std::function<int(Coroutine *)> &fun, const void *userData);

    virtual ~Coroutine();

    void resume();

    void yield();
};


#endif //HELLO_SWAPCONTEXT_COROUTINE_H
